from torch_geometric.datasets import OGB_MAG

dataset = OGB_MAG(root='/home/Dyf/code/dataset/data', preprocess='metapath2vec')
# preprocess也可以用TransE等
data = dataset[0]
print(data)
"""
我们来简单说明一下以下的数据
paper 论文
author  作者
institution 隶属机构
field_of_study  研究领域
HeteroData(
  paper={
    本数据集描述了 736389 篇论文和其他节点之间的纠缠关系。其中每一篇论文都用128个维度的向量去表示。
    x=[736389, 128],
    736389 篇论文分别发布于不同的年份。
    year=[736389],
    y 指的是每一篇论文都有一个类别标签。
    y=[736389],
    mask 说明针对这个数据集合我们将每一篇论文都安排到某一个数据集合中，train_mask中数据为 1 的位置上的论文隶属于训练集合，
    验证集和测试集按照相同的编码方式进行归纳操作。
    train_mask=[736389],
    val_mask=[736389],
    test_mask=[736389],
  },
  这些论文涉及的作者有 1134649 位，每一位都用128维的向量表示。
  author={ x=[1134649, 128] },
  这些论文涉及的机构有 8740 个，每一个都用128维的向量表示。
  institution={ x=[8740, 128] },
  这些论文涉及的领域有 59965 个，每一个领域都用128维的向量表示。
  field_of_study={ x=[59965, 128] },
  作者和隶属机构之间的关系有 1043998 条 ，描述为一个 [2, 1043998] 的矩阵形式，其中第一行向量表示 作者集合编号，第二行向量表示 机构集合编号。
  (author, affiliated_with, institution)={ edge_index=[2, 1043998] },
  作者写作论文是另外一种关系，这种关系有 7145660 条
  (author, writes, paper)={ edge_index=[2, 7145660] },
  论文引用论文是第三种关系，这种关系有 5416271 条
  (paper, cites, paper)={ edge_index=[2, 5416271] },
  论文和所涉及的研究领域是第四种关系，这种关系有 7505078 条
  (paper, has_topic, field_of_study)={ edge_index=[2, 7505078] }
)
"""

print(data.x_dict["paper"].shape)
print(data.x_dict["author"].shape)
print(data.x_dict["institution"].shape)
print(data.x_dict["field_of_study"].shape)
"""
以下是所有的x节点的信息
torch.Size([736389, 128])
torch.Size([1134649, 128])
torch.Size([8740, 128])
torch.Size([59965, 128])
"""

print(data.y_dict)
"""
这里只针对paper进行了分类
{'paper': tensor([246, 131, 189,  ..., 266, 289,   1])}
"""

print(data.year_dict)
"""
key 仍然是论文 value是一个年份组成的向量
{'paper': tensor([2015, 2012, 2012,  ..., 2016, 2017, 2014])}
"""


def count_sum(data):
    count = data.shape[0]
    count_true = data.sum().item()
    count_false = data.shape[0] - data.sum().item()
    # print(count_true)
    # print(count_false)
    percent = count_true / count
    return percent


print(count_sum(data.train_mask_dict["paper"]))
print(count_sum(data.val_mask_dict["paper"]))
print(count_sum(data.test_mask_dict["paper"]))
"""
对paper数据进行分割成的训练、验证、测试集合占比
训练 0.8549435149085606
验证 0.08810424924869872
测试 0.05695223584274073
{'paper': tensor([True, True, True,  ..., True, True, True])}
"""

"""
综合以上，我们发现一种反嵌套的结构，这个结构 {变量}_dict 中的变量扫描了所有数据结构中的所有字典信息，如果字典信息中存在该key，则将该信息调出 
例如 data.x_dict["paper"] 中 paper是key，paper中存在x的key，所以可以取到 data.x_dict["paper"]的数据。按照该逻辑推算，如果我们
取得 edge_index_dict信息，则会把所有有 edge_index信息的边的key都调出来，组成新的字典。下面的代码将验证这一点
"""
for key in data.edge_index_dict.keys():
    print(key)
    print(data.edge_index_dict[key].shape)
"""
('author', 'affiliated_with', 'institution')
torch.Size([2, 1043998])
('author', 'writes', 'paper')
torch.Size([2, 7145660])
('paper', 'cites', 'paper')
torch.Size([2, 5416271])
('paper', 'has_topic', 'field_of_study')
torch.Size([2, 7505078])
"""

"""
节点的两种表示法
"""
paper_node_data = data['paper']['x'].shape
paper_node_data1 = data['paper'].x.shape
"""
边的4种表示法
"""
cites_edge_data = data['paper', 'cites', 'paper'].edge_index.shape
cites_edge_data1 = data[('paper', 'cites', 'paper')].edge_index.shape
# 如果边类型或节点对类型可以唯一确定一种边，以下两种方式是支持的
cites_edge_data2 = data['paper', 'paper'].edge_index.shape
cites_edge_data3 = data['cites'].edge_index.shape
print(paper_node_data)
print(paper_node_data1)
print(cites_edge_data)
print(cites_edge_data1)
print(cites_edge_data2)
print(cites_edge_data3)
"""
总结，表示法方面相当丰富，采用其中一种方式即可。采用字典方式，方便后期进行异构计算。
"""

"""
这里只是一种示例，实际的作用是有些针对数据的变换可以后置进行，例如将所有x的128维度的向量做变换之后，得到的新的特征可以赋给原有data信息
"""
# 针对数据的修改
# data['paper'].year = data['paper'].y    # Setting a new paper attribute
# print(data['paper'].year)

"""
这里举例子，如果要去掉研究领域信息以及论文和研究领域之间的关闭，就可以使用删除key的方式来进行。这样得到的新的数据，将不会存再和研究领域有关的信息
"""
del data['field_of_study']  # Deleting 'field_of_study' node type
del data['has_topic']  # Deleting 'has_topic' edge type
print(data)

"""
这里获取到了data的元数据信息，包含节点信息和编信息的类型
"""
node_types, edge_types = data.metadata()
print(node_types)
print(edge_types)

# 判断是否有孤立节点
print(data.has_isolated_nodes())
# 判断是否有自环
print(data.has_self_loops())
# 判断是否无向图
print(data.is_undirected())
homogeneous_data = data.to_homogeneous()
print(homogeneous_data)
print(homogeneous_data.year)
print(homogeneous_data.y)
print(homogeneous_data.train_mask.sum().item())
print(data["paper"].train_mask.sum().item())
# 629571 + 1134649 + 8740 = 1772960 和 同构图中的train_mask 为True的值相同，这说明所有的没有指定train_mask的节点都是TRUE。
print(data["paper"].val_mask.sum().item() + 1134649 + 8740)
print(homogeneous_data.val_mask.sum().item())
print(data["paper"].test_mask.sum().item() + 1134649 + 8740)
print(homogeneous_data.test_mask.sum().item())
"""
HeteroData(
  paper={
    x=[736389, 128],
    year=[736389],
    y=[736389],
    train_mask=[736389],
    val_mask=[736389],
    test_mask=[736389],
  },
  author={ x=[1134649, 128] },
  institution={ x=[8740, 128] },
  (author, affiliated_with, institution)={ edge_index=[2, 1043998] },
  (author, writes, paper)={ edge_index=[2, 7145660] },
  (paper, cites, paper)={ edge_index=[2, 5416271] }
)

Data(
# 边的集合
edge_index=[2, 13605929], 
# 点的集合
x=[1879778, 128], 

tensor([2015, 2012, 2012,  ...,   -1,   -1,   -1]) 多余的值给 -1
year=[1879778], 
tensor([246, 131, 189,  ...,  -1,  -1,  -1])  多余的值给 -1
y=[1879778], 
train_mask=[1879778], 
val_mask=[1879778], 
test_mask=[1879778], 
node_type=[1879778], 
edge_type=[13605929])

"""

import torch_geometric.transforms as T

data = T.ToUndirected()(data)
data = T.AddSelfLoops()(data)
data = T.NormalizeFeatures()(data)
print(data)
"""
HeteroData(
  paper={
    x=[736389, 128],
    year=[736389],
    y=[736389],
    train_mask=[736389],
    val_mask=[736389],
    test_mask=[736389],
  },
  author={ x=[1134649, 128] },
  institution={ x=[8740, 128] },
  (author, affiliated_with, institution)={ edge_index=[2, 1043998] },
  (author, writes, paper)={ edge_index=[2, 7145660] },
  (paper, cites, paper)={ edge_index=[2, 5416271] }
)
HeteroData(
  paper={
    x=[736389, 128],
    year=[736389],
    y=[736389],
    train_mask=[736389],
    val_mask=[736389],
    test_mask=[736389],
  },
  author={ x=[1134649, 128] },
  institution={ x=[8740, 128] },
  (author, affiliated_with, institution)={ edge_index=[2, 1043998] },
  (author, writes, paper)={ edge_index=[2, 7145660] },
  (paper, cites, paper)={ edge_index=[2, 11529061] },
  (institution, rev_affiliated_with, author)={ edge_index=[2, 1043998] },
  (paper, rev_writes, author)={ edge_index=[2, 7145660] }
)
"""